# -*- coding:utf-8 -*-

# @Time    : 2023/8/15 07:05
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : memory.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# 用于存储用户对话记录中的信息，以便后续使用
import json


class Memory:
    def __init__(self) -> None:
        self.memory_cache = {}

    def set_user_info_memory(self, session_id, value):
        if session_id not in self.memory_cache.keys():
            self.memory_cache[session_id] = {}
        self.memory_cache[session_id]['user_info'] = value

    def get_user_info_memory(self, session_id):
        if session_id not in self.memory_cache.keys():
            return None
        if 'user_info' not in self.memory_cache[session_id].keys():
            return None
        return self.memory_cache[session_id]['user_info']

    def set_solution_memory(self, session_id, value):
        if session_id not in self.memory_cache.keys():
            self.memory_cache[session_id] = {}
        self.memory_cache[session_id]['solution'] = value

    def get_solution_memory(self, session_id):
        if session_id not in self.memory_cache.keys():
            return None
        if 'solution' not in self.memory_cache[session_id].keys():
            return None
        return self.memory_cache[session_id]['solution']

    def get_user_causal_graph(self, session_id):
        if session_id not in self.memory_cache.keys():
            return None
        if "user_causal_graph" not in self.memory_cache[session_id].keys():
            return None
        return self.memory_cache[session_id]['user_causal_graph']

    def set_user_causal_graph(self, session_id, user_causal_graph):
        if session_id not in self.memory_cache.keys():
            self.memory_cache[session_id] = {}
        self.memory_cache[session_id]['user_causal_graph'] = user_causal_graph

    def get_memory(self, hkey, key):
        if hkey in self.memory_cache and key in self.memory_cache[hkey]:
            value_json = self.memory_cache[hkey][key]
            if value_json:
                return json.loads(value_json)
        return None

    def set_memory(self, hkey, key, value):
        value_json = json.dumps(value, default=list, ensure_ascii=False)
        if hkey not in self.memory_cache:
            self.memory_cache[hkey] = {}
        self.memory_cache[hkey][key] = value_json

    def hexist_key(self, hkey, key):
        return hkey in self.memory_cache and key in self.memory_cache[hkey]

    def hdel(self, hkey, key):
        if hkey in self.memory_cache and key in self.memory_cache[hkey]:
            del self.memory_cache[hkey][key]

    def set(self, key, value):
        self.memory_cache[key] = value

    def get(self, key):
        return self.memory_cache.get(key)

    def keys(self, pattern):
        return [key for key in self.memory_cache.keys() if key.startswith(pattern)]

    def sadd(self, key, value):
        if key in self.memory_cache:
            self.memory_cache[key].add(value)
        else:
            self.memory_cache[key] = set([value])

    def smembers(self, key):
        return list(self.memory_cache.get(key, set()))

    def srem(self, key, value):
        if key in self.memory_cache and value in self.memory_cache[key]:
            self.memory_cache[key].remove(value)

    def delete(self, key):
        if key in self.memory_cache:
            del self.memory_cache[key]

    def hgetall(self, hkey):
        return self.memory_cache.get(hkey, {})

    def exist_key(self, key):
        return key in self.memory_cache.keys()
